In [1]:

    

import torch
import torchvision


    

In [2]:

    

torch.__version__


    

    
Out[2]:





'1.0.0.dev20181116'

In [3]:

    

# Let's create an example model using ResNet-18
model = torchvision.models.resnet18()


    

In [4]:

    

model


    

    
Out[4]:





ResNet(
  (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (relu): ReLU(inplace)
  (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
  (layer1): Sequential(
    (0): BasicBlock(
      (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace)
      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
    (1): BasicBlock(
      (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace)
      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
  )
  (layer2): Sequential(
    (0): BasicBlock(
      (conv1): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace)
      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (downsample): Sequential(
        (0): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)
        (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
    )
    (1): BasicBlock(
      (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace)
      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
  )
  (layer3): Sequential(
    (0): BasicBlock(
      (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace)
      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (downsample): Sequential(
        (0): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)
        (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
    )
    (1): BasicBlock(
      (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace)
      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
  )
  (layer4): Sequential(
    (0): BasicBlock(
      (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace)
      (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (downsample): Sequential(
        (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)
        (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
    )
    (1): BasicBlock(
      (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace)
      (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    )
  )
  (avgpool): AvgPool2d(kernel_size=7, stride=1, padding=0)
  (fc): Linear(in_features=512, out_features=1000, bias=True)
)

In [5]:

    

# Creating a sample of the input
# It will be used to pass it to the network to build the dimensions
sample = torch.rand(size=(1, 3, 224, 224))


    

In [6]:

    

# Creating so called "traced Torch script"
traced_script_module = torch.jit.trace(model, sample)


    

In [7]:

    

traced_script_module


    

    
Out[7]:





TracedModule[ResNet](
  (conv1): TracedModule[Conv2d]()
  (bn1): TracedModule[BatchNorm2d]()
  (relu): TracedModule[ReLU]()
  (maxpool): TracedModule[MaxPool2d]()
  (layer1): TracedModule[Sequential](
    (0): TracedModule[BasicBlock](
      (conv1): TracedModule[Conv2d]()
      (bn1): TracedModule[BatchNorm2d]()
      (relu): TracedModule[ReLU]()
      (conv2): TracedModule[Conv2d]()
      (bn2): TracedModule[BatchNorm2d]()
    )
    (1): TracedModule[BasicBlock](
      (conv1): TracedModule[Conv2d]()
      (bn1): TracedModule[BatchNorm2d]()
      (relu): TracedModule[ReLU]()
      (conv2): TracedModule[Conv2d]()
      (bn2): TracedModule[BatchNorm2d]()
    )
  )
  (layer2): TracedModule[Sequential](
    (0): TracedModule[BasicBlock](
      (conv1): TracedModule[Conv2d]()
      (bn1): TracedModule[BatchNorm2d]()
      (relu): TracedModule[ReLU]()
      (conv2): TracedModule[Conv2d]()
      (bn2): TracedModule[BatchNorm2d]()
      (downsample): TracedModule[Sequential](
        (0): TracedModule[Conv2d]()
        (1): TracedModule[BatchNorm2d]()
      )
    )
    (1): TracedModule[BasicBlock](
      (conv1): TracedModule[Conv2d]()
      (bn1): TracedModule[BatchNorm2d]()
      (relu): TracedModule[ReLU]()
      (conv2): TracedModule[Conv2d]()
      (bn2): TracedModule[BatchNorm2d]()
    )
  )
  (layer3): TracedModule[Sequential](
    (0): TracedModule[BasicBlock](
      (conv1): TracedModule[Conv2d]()
      (bn1): TracedModule[BatchNorm2d]()
      (relu): TracedModule[ReLU]()
      (conv2): TracedModule[Conv2d]()
      (bn2): TracedModule[BatchNorm2d]()
      (downsample): TracedModule[Sequential](
        (0): TracedModule[Conv2d]()
        (1): TracedModule[BatchNorm2d]()
      )
    )
    (1): TracedModule[BasicBlock](
      (conv1): TracedModule[Conv2d]()
      (bn1): TracedModule[BatchNorm2d]()
      (relu): TracedModule[ReLU]()
      (conv2): TracedModule[Conv2d]()
      (bn2): TracedModule[BatchNorm2d]()
    )
  )
  (layer4): TracedModule[Sequential](
    (0): TracedModule[BasicBlock](
      (conv1): TracedModule[Conv2d]()
      (bn1): TracedModule[BatchNorm2d]()
      (relu): TracedModule[ReLU]()
      (conv2): TracedModule[Conv2d]()
      (bn2): TracedModule[BatchNorm2d]()
      (downsample): TracedModule[Sequential](
        (0): TracedModule[Conv2d]()
        (1): TracedModule[BatchNorm2d]()
      )
    )
    (1): TracedModule[BasicBlock](
      (conv1): TracedModule[Conv2d]()
      (bn1): TracedModule[BatchNorm2d]()
      (relu): TracedModule[ReLU]()
      (conv2): TracedModule[Conv2d]()
      (bn2): TracedModule[BatchNorm2d]()
    )
  )
  (avgpool): TracedModule[AvgPool2d]()
  (fc): TracedModule[Linear]()
)

In [8]:

    

# The TracedModule is capable of making predictions
sample_prediction = traced_script_module(torch.ones(size=(1, 3, 224, 224)))


    

In [9]:

    

sample_prediction.shape


    

    
Out[9]:





torch.Size([1, 1000])

In [11]:

    

# Serializing the the script module
traced_script_module.save('./models_deployment/model.pt')